(*  CSIsat: interpolation procedure for LA + EUF
 *  Copyright (C) 2008  The CSIsat team
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *)

open Ast
open AstUtil
open DpllClause
open DpllProof

(*a simple dpll SAT solver*)
(*TODO improve*)

type var_assign = Open
                | Closed of res_proof
                | Implied of res_proof
                | TImplied of res_proof

class system =
  fun with_prf ->
  object (self)
    
    val mutable possibly_sat = true
    val mutable resolution_proof = None
    
    val mutable clauses = Array.make 0 (new clause (Or [Eq(Constant 1.0,Constant 1.0)]) true )
    val mutable props = PredSet.empty
    val mutable affected = PredSet.empty
    val choices = Stack.create ()

    (*2-list watching*)
    val prop_to_clauses = Hashtbl.create 123
    
    val keep_proof = with_prf
    val mutable learning_level = 1
    (** -1: no learning
     *  0: some learning
     *  1+: learn clause that are less or equal than value
     * Warning, cal this method after having called init
     *)
    method set_learning_level l = learning_level <- l

    (*is an OrdSet*)
    val mutable unsat_clauses = []

    method reset =
      clauses <- Array.make 0 (new clause (Or [Eq(Constant 1.0,Constant 1.0)]) true);
      unsat_clauses <- [];
      props <- PredSet.empty;
      affected <- PredSet.empty;
      Hashtbl.clear prop_to_clauses;
      Stack.clear choices

    (*2-list watching as ordset*)
    method add_pos_clause_for_prop p cl =
      let (oldp, oldn) = try Hashtbl.find prop_to_clauses p with Not_found -> ([],[]) in
        Hashtbl.replace prop_to_clauses p (OrdSet.union [cl] oldp, oldn)
    method add_neg_clause_for_prop p cl =
      let (oldp, oldn) = try Hashtbl.find prop_to_clauses p with Not_found -> ([],[]) in
        Hashtbl.replace prop_to_clauses p (oldp, OrdSet.union [cl] oldn)

    
    (*assume CNF*)
    method init formula = 
      props <- get_proposition_set formula;
      PredSet.iter (fun x -> Hashtbl.add prop_to_clauses x ([],[])) props;
      match formula with
      | And lst ->
        begin
          let n = List.length lst in
            clauses <- Array.make n (new clause (Or [Eq(Constant 1.0,Constant 1.0)]) true);
            ignore (List.fold_left (fun i e -> 
                let cl = new clause e false in
                  clauses.(i) <- cl;
                  let pos = cl#pos_props in PredSet.iter (fun x -> self#add_pos_clause_for_prop x cl) pos;
                  let neg = cl#neg_props in PredSet.iter (fun x -> self#add_neg_clause_for_prop x cl) neg;
                    i + 1
              ) 0 lst);
            unsat_clauses <- OrdSet.list_to_ordSet (Array.to_list clauses);
            let average_size = (Array.fold_left (fun acc cl -> acc + cl#size) 0 clauses) / (Array.length clauses) in
              self#set_learning_level ((3 * average_size) / 2)
        end
      | _ -> failwith "DPLL: expect CNF"
     
    method has_contra = 
      List.exists (fun x -> x#contradiction) unsat_clauses

    method is_sat =
      match unsat_clauses with
      | [] -> true
      | _ -> false

    method affect p reason =
      Message.print Message.Debug (lazy("DPLL, affecting : "^(print p)));
      assert (not (PredSet.mem (contra p) affected));
      affected <- PredSet.add p affected;
      let (pos,neg) = Hashtbl.find prop_to_clauses (List.hd (get_proposition p)) in
      let (_true,_false) = if (List.hd (get_proposition p)) = p then (pos,neg) else (neg,pos)
      in
      let newly_sat = List.filter (fun x -> not x#is_sat) _true in
        List.iter (fun x -> x#affect p) _false;
        List.iter (fun x -> x#affect p) _true;
        unsat_clauses <- OrdSet.substract unsat_clauses newly_sat;
        Stack.push (p,reason, newly_sat) choices

    method forget =
      let (pivot,how,satisfied) = Stack.pop choices in
      Message.print Message.Debug (lazy("DPLL, forgetting: "^(print pivot)));
      assert (PredSet.mem pivot affected);
      affected <- PredSet.remove pivot affected;
      let (pos,neg) = Hashtbl.find prop_to_clauses (List.hd (get_proposition pivot)) in
      let (_true,_false) = if (List.hd (get_proposition pivot)) = pivot then (pos,neg) else (neg,pos)
      in
        List.iter (fun x -> x#forget pivot) _true;
        List.iter (fun x -> x#forget pivot) _false;
        unsat_clauses <- OrdSet.union unsat_clauses satisfied;
        (pivot,how)

    method get_assign =
      PredSet.fold (fun a acc -> a::acc) affected []
    
    method get_assigned_props =
      List.flatten (List.map get_proposition self#get_assign)

    (*return the first clause that satisfies fct*)
    method scan_clauses fct =
      let i = ref 0 in
      let n = Array.length clauses in
      let ans = ref false in
        while (not !ans) && !i < n do
          ans := fct clauses.(!i);
          i := !i + 1
        done;
        if !ans then Some clauses.(!i-1) else None
    
    method filter_clauses fct =
      List.filter fct (Array.to_list clauses)
    
    method iter_clauses fct =
      Array.iter fct clauses

    method to_string =
      let buffer = Buffer.create 1024 in
      let assign = self#get_assign in
        Buffer.add_string buffer "current assign is ";
        List.iter (fun x ->
            Buffer.add_string buffer (print x); 
            Buffer.add_string buffer ", ";
          ) assign;
        Buffer.add_char buffer '\n';
        self#iter_clauses (fun x ->
          Buffer.add_string buffer (x#to_string_detailed);
          Buffer.add_char buffer '\n');
        Buffer.contents buffer

    (************************************************************)
    (****************   DECISION POLICY   ***********************)
    (************************************************************)

    (*return a variable that only satisfies clauses *)
    method find_max_unipolar_variable =
      let max = ref 0 in
      let prop = ref None in
        Hashtbl.iter
          (fun pr (pos,neg) ->
            if not (PredSet.mem pr affected) && not (PredSet.mem (contra pr) affected) then
              begin
                let pos = List.filter (fun x -> not x#is_sat) pos in
                let neg = List.filter (fun x -> not x#is_sat) neg in
                  match (pos,neg) with
                  | ([],[]) -> ()
                  | (p, []) -> if List.length p > !max then (prop := Some pr; max := List.length p)
                  | ([], n) -> if List.length n > !max then (prop := Some (contra pr); max := List.length n)
                  | _ -> ()
              end
          ) prop_to_clauses;
        !prop

    (** try to find a clause with only one literal left.
     * @return Some(literal,clause)
     *)
    method find_unit_propagation =
      try 
        let p = List.find (fun x -> x#size = 1) unsat_clauses in
        let c = p#get_choice in
          Some (c,p)
      with Not_found -> None
      (*match self#scan_clauses (fun x -> x#size = 1 && not x#is_sat) with
      | Some p -> let c = p#get_choice in Some (c,p)
      | None -> None*)

    (*return a variable that satisfies the maximun #clauses *)
    method find_max_degree =
      let max = ref 0 in
      let prop = ref None in
        Hashtbl.iter
          (fun pr (pos,neg) ->
            if not (PredSet.mem pr affected) && not (PredSet.mem (contra pr) affected) then
              begin
                let pos = List.filter (fun x -> not x#is_sat) pos in
                let neg = List.filter (fun x -> not x#is_sat) neg in
                  if abs ((List.length pos) - (List.length neg)) > !max then
                    begin
                      max := abs ((List.length pos) - (List.length neg));
                      if (List.length pos) - (List.length neg) > 0
                        then prop := Some pr
                        else prop := Some (contra pr)
                    end
              end
          ) prop_to_clauses;
        !prop
      
    (** return a literal that will make a clause sat*)
    method find_random =
      let length = List.length unsat_clauses in
      let n = Random.int length in
      let c = (List.nth unsat_clauses n) in
        Some c#get_choice
      

    (** return a clause that cannot be satisfied*)
    method get_unsat_clause =
      try List.find (fun x -> x#contradiction) unsat_clauses
      with Not_found -> failwith "DPLL: get_unsat_clause called without contradiction"

    (** insert a new clause
     * @return false if need to backtrack
     *)
    method new_clause cl =
      let pos = cl#pos_props in PredSet.iter (fun x -> self#add_pos_clause_for_prop x cl) pos;
      let neg = cl#neg_props in PredSet.iter (fun x -> self#add_neg_clause_for_prop x cl) neg;
      List.iter (cl#affect) self#get_assign;
      clauses <- Array.append clauses (Array.make 1 cl);
      if cl#is_sat then
        begin
          let sat_element = ref (cl#get_satisfied) in
          (*unstack to the first lit that satisfied the clause*)
          let copy = ref [] in
            while (PredSet.cardinal !sat_element) > 0 do
              let (pivot,_,_) as entry = Stack.pop choices in
                copy := entry::!copy;
                sat_element := PredSet.remove pivot !sat_element
            done;
            let (pivot,reason,clauses) = List.hd !copy in
              assert (cl#has pivot);
              Stack.push (pivot,reason, OrdSet.union [cl] clauses) choices;
              List.fold_left (fun () x -> Stack.push x choices) () (List.tl !copy);
              true
        end
      else
        begin
          unsat_clauses <- OrdSet.union [cl] unsat_clauses;
          not (cl#contradiction)
        end

    method learned_clause disj =
      let cl = new clause disj true in
      let res = self#new_clause cl in
        assert res
    
    method learn_clause cl =
      ignore (self#new_clause cl)
      (*
      let res = self#new_clause cl in
        assert res
      *)
    
    val partial_proof = Hashtbl.create 1000
    method store_proof clause proof = Hashtbl.replace partial_proof clause proof
    method get_partial_proof clause = Hashtbl.find partial_proof clause

    
    (** to call when unsat + need a proof
     *)
    method backjump explanation =
      let clause_of_set set =
        new clause (Or (PredSet.fold (fun x acc -> x::acc) set [])) true
      in
      let rec build_proof prf =
        try 
          let (pivot, how) = self#forget in
            match how with
            | Open -> (*can stop the proof here and pick new assign*)
              begin
                let assigned = OrdSet.list_to_ordSet (self#get_assigned_props) in
                let choices = OrdSet.list_to_ordSet (PredSet.fold (fun x acc -> x::acc) ((get_result prf)#props) []) in
                let resulting = OrdSet.substract choices assigned in
                  if List.length resulting > 0 then
                    begin
                      Message.print Message.Debug (lazy("DPLL, backjumping ended in Open"));
                      let c =
                        let tmp = (List.hd resulting) in
                          if (get_result prf)#has tmp then
                            tmp
                          else
                            begin
                              assert ((get_result prf)#has_not tmp);
                              contra tmp
                            end

                      in
                        if keep_proof then
                          self#affect c (Closed prf)
                        else
                          self#affect c (Closed (RPLeaf (get_result prf)))
                    end
                  else
                    build_proof prf
              end
            | Closed proof -> (*try to find an not tested var or continue proof further*)
              begin
                let satisfied_clause = get_result proof in
                let to_satisfy = get_result prf in
                let new_prf =
                  if to_satisfy#has_not pivot then
                    begin
                      let new_unsat_disj = 
                        clause_of_set
                          (PredSet.remove (contra pivot)
                            (PredSet.remove pivot
                              (PredSet.union ((get_result prf)#get_propositions) (satisfied_clause#get_propositions))))
                      in
                        assert(not new_unsat_disj#has_el_and_contra);
                        let new_prf = if keep_proof then
                            RPNode (List.hd (get_proposition pivot), proof, prf, new_unsat_disj)
                          else
                            RPLeaf new_unsat_disj
                        in
                          if learning_level >= 1 && new_unsat_disj#size <= learning_level then
                            begin
                              self#store_proof new_unsat_disj new_prf;
                              self#learn_clause new_unsat_disj
                            end;
                          new_prf
                    end
                  else
                    begin
                      (*at this point, it may be usefull to learn the clause of (Closed proof)*)
                      (*also keep the proof*)
                      if learning_level = 0 then
                        begin
                          self#store_proof satisfied_clause proof;
                          self#learn_clause satisfied_clause
                        end;
                      prf
                    end
                in
                  let possible = PredSet.fold (fun x acc -> x::acc ) ((get_result new_prf)#get_propositions) [] in
                    try
                      let new_try = List.find
                        (fun x ->
                          (not (PredSet.mem x affected)) && (not (PredSet.mem (contra x) affected))
                        ) possible
                      in
                        Message.print Message.Debug (lazy("DPLL, backjumping ended in Closed"));
                        self#affect new_try (Closed new_prf)
                    with Not_found ->
                      build_proof new_prf
              end
            | Implied proof -> (*continue proof further*)
              begin
                let satisfied_clause = get_result proof in
                let to_satisfy = get_result prf in
                let new_prf =
                  if to_satisfy#has_not pivot then
                    begin
                      let new_unsat_disj = 
                        clause_of_set
                          (PredSet.remove (contra pivot)
                            (PredSet.remove pivot
                              (PredSet.union ((get_result prf)#get_propositions) (satisfied_clause#get_propositions))))
                      in
                        assert(not new_unsat_disj#has_el_and_contra);
                        let new_prf = if keep_proof then
                            RPNode (List.hd (get_proposition pivot), proof, prf, new_unsat_disj)
                          else
                            RPLeaf new_unsat_disj
                        in
                          if learning_level >= 1 && new_unsat_disj#size <= learning_level then
                            begin
                              self#store_proof new_unsat_disj new_prf;
                              self#learn_clause new_unsat_disj
                            end;
                          new_prf
                    end
                  else
                    begin
                      prf
                    end
                in
                  build_proof new_prf
              end
            | TImplied proof -> (*continue proof further*)
              begin
                (*TODO need T-lemma*)
                failwith "DPLL: theory deduction yet to come"
              end
        with Stack.Empty ->
          begin (*now we have a proof of unsat*)
            Message.print Message.Debug (lazy(tracecheck_of_proof prf));
            assert ((get_result prf)#get_propositions = PredSet.empty);
            resolution_proof <- Some prf;
            possibly_sat <- false
          end
      in
      let start_proof = 
        if explanation#is_learned then
          self#get_partial_proof explanation
        else
          RPLeaf explanation
      in
        build_proof start_proof


    method add_clause disj =
      let cl = new clause disj false in
      let res = self#new_clause cl in
        if not res then self#backjump cl
   
    (*TODO T-propagation*)
    method decision_policy =
      match self#find_unit_propagation with
      | Some (lit,cl) ->
        begin
          Message.print Message.Debug (lazy("DPLL, found unit propagation: "^(print lit)));
          self#affect (lit) (Implied (RPLeaf cl))
        end
      | None ->
        begin
          Message.print Message.Debug (lazy("DPLL, no unit propagation"));
          match self#find_max_unipolar_variable with
          | Some lit  ->
            begin
              Message.print Message.Debug (lazy("DPLL, found max unipolar literal: "^(print lit)));
              self#affect lit Open
            end
          | None ->
            begin
              Message.print Message.Debug (lazy("DPLL, no max degree literal"));
              match self#find_max_degree with
              | Some lit  ->
                begin
                  Message.print Message.Debug (lazy("DPLL, found max max degree literal: "^(print lit)));
                  self#affect lit Open
                end
              | None ->
                begin
                  match self#find_random with
                  | Some lit  ->
                    begin
                      Message.print Message.Debug (lazy("DPLL, found random literal: "^(print lit)));
                      self#affect lit Open
                    end
                  | None ->
                    failwith "DPLL, decision_policy: no possible affectation"
                end
            end
        end

    method solve =
      (*Message.print Message.Debug (lazy("DPLL, status: "^(self#to_string)));*)
      if possibly_sat then
        begin
        Message.print Message.Debug (lazy("DPLL, system is possibly sat."));
          if self#is_sat then
            Some(self#get_assign)
          else if self#has_contra then
            begin
              let cl = self#get_unsat_clause in
                Message.print Message.Debug (lazy("DPLL, backtracking with: "^(cl#to_string)));
                self#backjump cl;
                self#solve
            end
          else
            begin
              Message.print Message.Debug (lazy("DPLL taking decision_policy branch"));
              self#decision_policy;
              self#solve
            end
        end
      else
        None

    method get_proof_of_unsat =
      match resolution_proof with
      | Some prf -> prf
      | None -> failwith "DPLL, no resolution proof"

  end

(*** Wrapper ***)
open SatInterface

class my_dpll with_proof =
  object
    inherit sat_solver with_proof
    val sys = new system with_proof

    method init formulae = sys#init formulae
    
    method add_clause formula = sys#add_clause formula
    
    val mutable last_solution: Ast.predicate list option = None
    method solve = match sys#solve with
      | Some sol -> last_solution <- Some sol; true
      | None -> last_solution <- None; false

    method get_solution = match last_solution with
      | Some sol -> List.map  normalize_only sol
      | None -> failwith "DpllCore: no solution for the moment"
    
    method get_proof = sys#get_proof_of_unsat
  end
